from langchain.chains import LLMChain
from langchain_community.chat_models import ChatTongyi
from langchain.prompts import ChatPromptTemplate, FewShotChatMessagePromptTemplate
import os


os.environ["DASHSCOPE_API_KEY"] = "sk-9d8f1914800e497f8717144e860f99bc"
llm = ChatTongyi()

# 生成具体回复
examples = [
    {"input": "2+2", "output": "4"},
    {"input": "2+3", "output": "5"},
]
example_prompt = ChatPromptTemplate.from_messages(
    [
        ("human", "{input}"),
        ("ai", "{output}")
    ]
)
few_shot_prompt = FewShotChatMessagePromptTemplate(
    example_prompt=example_prompt,
    examples=examples,
)

final_prompt = ChatPromptTemplate.from_messages(
                [
                    ('system', 'You are a helpful AI Assistant'),
                    few_shot_prompt,
                    ('human', '{input}'),
                ]
            )
final_prompt.format(input="What is 4+4?")

# 创建执行链
chain = LLMChain(
    llm=llm,
    prompt=final_prompt
)
res = chain.invoke({"input": "2+2 等于多少"})
res2 = chain.invoke({"input": "2+6 等于多少"})
print(res)
print(res2)


